# -*- coding: utf-8 -*-
"""
数据库配置模块

管理 PostgreSQL + pgvector 数据库的连接配置和相关设置。
支持向量存储和传统关系型数据的混合存储方案。
"""
import os
from typing import Optional
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from loguru import logger


class DatabaseConfig:
    """
    数据库配置类
    
    管理 PostgreSQL 数据库连接参数和会话配置。
    支持 pgvector 扩展，用于向量数据存储和检索。
    """
    
    # ==================== 数据库连接配置 ====================
    DB_HOST: str = os.getenv("DB_HOST", "10.48.0.81")
    DB_PORT: int = int(os.getenv("DB_PORT", "54333"))
    DB_NAME: str = os.getenv("DB_NAME", "langchain_pgvector")
    DB_USER: str = os.getenv("DB_USER", "pgvector")
    DB_PASSWORD: str = os.getenv("DB_PASSWORD", "pgvector")
    
    # ==================== 连接池配置 ====================
    # 连接池大小配置
    POOL_SIZE: int = int(os.getenv("DB_POOL_SIZE", "10"))
    MAX_OVERFLOW: int = int(os.getenv("DB_MAX_OVERFLOW", "20"))
    POOL_TIMEOUT: int = int(os.getenv("DB_POOL_TIMEOUT", "30"))
    POOL_RECYCLE: int = int(os.getenv("DB_POOL_RECYCLE", "3600"))
    
    # ==================== 向量配置 ====================
    # 向量维度（与嵌入模型保持一致）
    VECTOR_DIMENSION: int = int(os.getenv("VECTOR_DIMENSION", "768"))
    
    @classmethod
    def get_database_url(cls) -> str:
        """
        构建数据库连接URL
        
        Returns:
            str: PostgreSQL 连接字符串
        """
        return f"postgresql://{cls.DB_USER}:{cls.DB_PASSWORD}@{cls.DB_HOST}:{cls.DB_PORT}/{cls.DB_NAME}"
    
    @classmethod
    def create_engine(cls):
        """
        创建数据库引擎
        
        Returns:
            Engine: SQLAlchemy 数据库引擎
        """
        engine = create_engine(
            cls.get_database_url(),
            pool_size=cls.POOL_SIZE,
            max_overflow=cls.MAX_OVERFLOW,
            pool_timeout=cls.POOL_TIMEOUT,
            pool_recycle=cls.POOL_RECYCLE,
            echo=False  # 生产环境设为 False
        )
        logger.info(f"数据库引擎创建成功: {cls.DB_HOST}:{cls.DB_PORT}/{cls.DB_NAME}")
        return engine
    
    @classmethod
    def create_session_factory(cls):
        """
        创建会话工厂
        
        Returns:
            sessionmaker: SQLAlchemy 会话工厂
        """
        engine = cls.create_engine()
        return sessionmaker(bind=engine)
    
    @classmethod
    def validate_connection(cls) -> bool:
        """
        验证数据库连接
        
        Returns:
            bool: 连接是否成功
        """
        try:
            engine = cls.create_engine()
            with engine.connect() as conn:
                # 检查 pgvector 扩展
                result = conn.execute(text("SELECT 1"))
                result.fetchone()
                logger.info("数据库连接验证成功")
                return True
        except Exception as e:
            logger.error(f"数据库连接验证失败: {e}")
            return False